import uuid
from datetime import datetime, timedelta

import pandas as pd

from env_config import ENV
from task_service import AbsTaskService
from x_yc import prediction_plan0822
from x_yc.db.in_service import ForecastConfig, get_forecast_config, WarehouseType, get_deliver_plan, get_last_days_data


class XYCTaskService(AbsTaskService):
    """
    预测
    """
    SPLIT_KEY = '\u0001'
    DATE_FORMAT = '%Y-%m-%d'
    PRE_DAYS = 90

    DEFAULT_VAL = 0

    name: str = 'YC'

    predict_date: datetime = None
    # 预测算法 参数配置
    forecast_config: ForecastConfig = None

    date_list: list[str] = []
    keys_list: list[str] = []

    f_date_list: list[str] = []

    def __init__(self):
        self.date_list: list[str] = []
        self.keys_list: list[str] = []
        self.f_date_list: list[str] = []

    def _pre_process(self):
        """
        数据前置处理
        将入参转换为 算法所需格式数据
        :return:
        """
        self.context.app.logger.info('[YC] task:%s, pre process begin', self.context.task_id)
        self.forecast_config = get_forecast_config()
        self.context.app.logger.info(f'[YC] 算法配置信息: {self.forecast_config}')
        # 获取时间数据
        self.predict_date = datetime.now()
        predict_date_str = self.context.predict_date
        if predict_date_str:
            self.predict_date = datetime.strptime(predict_date_str, self.DATE_FORMAT)
        # 向前推 fix_data_range (考虑数据范围) 天的日期
        for idx in range(self.forecast_config.fix_data_range):
            date = self.predict_date + timedelta(days=-idx - 1)
            self.date_list.append(date.strftime(self.DATE_FORMAT))
        self.date_list = sorted(set(self.date_list), key=lambda _: datetime.strptime(_, self.DATE_FORMAT))

        # 向后推进 forecast_day (最大预测天数) 的日期
        for idx in range(self.forecast_config.fix_forecast_day):
            date = self.predict_date + timedelta(days=idx)
            self.f_date_list.append(date.strftime(self.DATE_FORMAT))
        self.f_date_list = sorted(set(self.f_date_list), key=lambda _: datetime.strptime(_, self.DATE_FORMAT))

        # 读取数据
        resp_data_list = self.context.in_data
        if not resp_data_list:
            resp_data_list = self.get_data(self.forecast_config.fix_data_range)

        # 转换
        in_data_map = {}
        for d in resp_data_list:
            if not d:
                continue
            key = self.SPLIT_KEY.join([d['type'], d['packageType'], d['inOutType'], str(d['warehouseType'])])
            if key not in in_data_map:
                in_data_map[key] = {}
            if d['date'] not in in_data_map[key]:
                in_data_map[key][d['date']] = 0
            in_data_map[key][d['date']] += float(d['totalNum'])
        # 按日期整理数据
        fixed_in_data_map = {}
        for key in in_data_map:
            fixed_in_data_map[key] = [in_data_map[key][_] if _ in in_data_map[key] else self.DEFAULT_VAL for _ in
                                      self.date_list]
        # 转为DataFrame
        self.df = pd.DataFrame(fixed_in_data_map)
        # 统计历史SKU比例
        self.h_plan_data = self.f_plan_data = pd.DataFrame({})
        self.context.app.logger.info('[YC] task:%s, consider_plan: {}', self.forecast_config.fix_consider_plan)
        if self.forecast_config.fix_consider_plan:
            type_nums_map = {}
            for d in resp_data_list:
                if not d:
                    continue
                key = d['type']  # sku
                if key not in type_nums_map:
                    type_nums_map[key] = {}
                    type_nums_map[key][WarehouseType.ZB.value] = 0
                    type_nums_map[key][WarehouseType.JX.value] = 0
                type_nums_map[key][str(d['warehouseType'])] += int(d['totalNum'])
            # 获取 (h历史) 计划数据
            hdp_list = get_deliver_plan(self.forecast_config.fix_data_range, history=True)
            h_plan_data_dict = self.process_feature_data(type_nums_map, hdp_list, self.date_list)
            # 获取 (f未来) 计划数据
            fdp_list = get_deliver_plan(self.forecast_config.fix_forecast_day, history=False)
            f_plan_data_dict = self.process_feature_data(type_nums_map, fdp_list, self.f_date_list)
            # 整理 聚合数据
            full_keys = set(fixed_in_data_map.keys()).union(h_plan_data_dict.keys()).union(f_plan_data_dict.keys())
            for key in full_keys:
                if key not in fixed_in_data_map:
                    fixed_in_data_map[key] = [0 for _ in range(self.forecast_config.fix_data_range)]
                if key not in h_plan_data_dict:
                    h_plan_data_dict[key] = [0 for _ in range(self.forecast_config.fix_data_range)]
                if key not in f_plan_data_dict:
                    f_plan_data_dict[key] = [0 for _ in range(self.forecast_config.fix_forecast_day)]
            # 最终生成数据
            self.df = pd.DataFrame(self.sorted_map(fixed_in_data_map))
            self.h_plan_data = pd.DataFrame(self.sorted_map(h_plan_data_dict))
            self.f_plan_data = pd.DataFrame(self.sorted_map(f_plan_data_dict))
        #
        self.keys_list = self.df.keys().tolist()
        self.context.app.logger.debug('[YC] task:%s, pre process data_map: %s', self.context.task_id, fixed_in_data_map)
        self.context.app.logger.info('[YC] task:%s, pre process end', self.context.task_id)

    def _process(self):
        """
        调用算法进行预测
        :return:
        """
        self.context.app.logger.info('[YC] task:%s, process begin', self.context.task_id)
        forecast_data = None
        if not self.df.empty:
            self.context.app.logger.debug('[YC] task:%s, process df: %s', self.context.task_id, self.df)
            self.context.app.logger.info('[YC] task:%s, process GPU-env: %s', self.context.task_id, ENV.gpu)
            try:
                self.context.app.logger.info(f'[YC] [{self.context.task_id}] 算法调用 预测 开始')
                forecast_data = prediction_plan0822.forecast(self.df.values, self.h_plan_data.values,
                                                             self.f_plan_data.values,
                                                             forecast_steps=self.forecast_config.fix_forecast_day,
                                                             use_plan=self.forecast_config.fix_consider_plan_sb,
                                                             use_gpu=ENV.gpu)
            except Exception as ex:
                self.context.app.logger.error(f'[YC] [{self.context.task_id}] 算法调用 预测 算法执行异常, ex: {ex}')
        else:
            self.context.app.logger.warning('[YC] task:%s, process error, no data', self.context.task_id)
        self.result = forecast_data
        self.context.app.logger.info('[YC] task:%s, process end', self.context.task_id)

    def _post_process(self):
        """
        数据后处理
        :return:
        """
        self.context.app.logger.info('[YC] task:%s, post process begin', self.context.task_id)
        super()._post_process()
        result_list = self.result.tolist() if self.result is not None else []
        #
        out_data = []
        for idx, row in enumerate(result_list):
            date_str = (self.predict_date + timedelta(days=(idx + 1))).strftime(self.DATE_FORMAT)
            for kidx, key in enumerate(self.keys_list):
                _type, _packageType, _inOutType, _warehouseType = key.split(self.SPLIT_KEY)
                out_data.append({'id': str(uuid.uuid1().hex), 'date': date_str, 'group': '', 'type': _type,
                                 'packageType': _packageType, 'inOutType': _inOutType, 'warehouseType': _warehouseType,
                                 'totalNum': self.fix_result_value(row[kidx]),
                                 'addTime': self.predict_date.strftime(self.DATE_FORMAT)})
        # 结果赋值
        self.context.out_data = out_data
        self.context.app.logger.info('[YC] task:%s, post process end', self.context.task_id)

    def process_feature_data(self, type_nums_map: dict, dp_list: list, date_list: list):
        dp_map = {}
        for dp in dp_list:
            total_num = int(dp.totalNum)
            if dp.type in type_nums_map:
                zb_num = type_nums_map[dp.type][WarehouseType.ZB.value]
                jx_num = type_nums_map[dp.type][WarehouseType.JX.value]
                # 平分秋色
                if zb_num == 0 and jx_num == 0:
                    zb_total = int(total_num / 2)
                elif zb_num == 0 and jx_num > 0:
                    zb_total = 0
                elif zb_num > 0 and jx_num == 0:
                    zb_total = total_num
                else:
                    zb_total = int(total_num / (zb_num + jx_num) * jx_num)
            else:
                zb_total = int(total_num / 2)
            jx_total = total_num - zb_total
            #
            date = dp.date
            zb_key = self.SPLIT_KEY.join([dp.type, dp.packageType, str(dp.inOutType), WarehouseType.ZB.value])
            jx_key = self.SPLIT_KEY.join([dp.type, dp.packageType, str(dp.inOutType), WarehouseType.JX.value])
            for k in (zb_key, jx_key):
                if k not in dp_map:
                    dp_map[k] = {}
                if date not in dp_map[k]:
                    dp_map[k][date] = 0
            dp_map[zb_key][date] += zb_total
            dp_map[jx_key][date] += jx_total
        plan_data_dict = {}
        for key in dp_map:
            plan_data_dict[key] = [dp_map[key][_] if _ in dp_map[key] else self.DEFAULT_VAL for _ in date_list]
        return plan_data_dict

    @staticmethod
    def fix_result_value(val: float):
        """
        对计算结果进行四舍五入
        :param val:
        :return:
        """
        if val < 0:
            val = 0.0
        return int(val + 0.5)

    @staticmethod
    def get_data(days: int) -> list:
        """
        根据参数 查询数据
        :return:
        """
        data_list = get_last_days_data(days)
        return [obj.to_dict() for obj in data_list]

    @staticmethod
    def sorted_map(data: dict) -> dict:
        """
        返回一个按Key值有序的?字典  供DataFrame使用
        """
        return dict(sorted(data.items(), key=lambda _: _[0]))
